# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/29 14:17 
@Author  : ZhangShenao 
@File    : 3.使用MMR最大边际相关性搜索.py
@Desc    : 使用MMR最大边际相关性搜索

最大边际相关性(MMR,max_marginal_relevance_search)的
基本思想是同时考量查询与文档的相关度,以及文档之间的相似度
在满足检索结果的相关性的前提下,尽量保证结果的多样性
它计算每个候选文档与查询的相关度,并减去与已经入选结果集的文档的最大相似度,这样更不相似的文档会有更高分

底层计算得=lambda_mult *相关性 - (1 - lambda_mult)*相似性
0代表最大多样性
1代表最小多样性

MMR的执行流程为:
1. 搜索出fetch_k条相关文档
2. 计算所有文档与query之间的相似性得分,和每个文档之间的相似性得分
3. 根据lambda_mult系数计算,综合计算上述两个相似性得分,得到最终的得分
4. 按照最终得分排序,返回最相关的k条结果

使用MMR检索,可以有效避免检索结果中出现重复的文档或相似度很高的文档的情况

"""

import dotenv
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_core.runnables import ConfigurableField
from langchain_pinecone import PineconeVectorStore

# 加载环境变量
dotenv.load_dotenv()

embeddings = ZhipuAIEmbeddings(model='embedding-2')
vector_store = PineconeVectorStore(
    index_name='dataset',  # 指定索引名称
    embedding=embeddings,  # 指定Embedding模型
    namespace='llm-ops',  # 指定Namespace
    text_key='text',  # 指定文本属性key
)

# Retrieving文档检索
# retriever是Runnable可运行组件,可以通过LCEL表达式直接运行
retriever = vector_store.as_retriever(
    search_type='mmr',  # 指定检索类型,使用MMR搜索,保证检索结果的多样性
    # 指定检索参数
    search_kwargs={
        'k': 3,  # 指定返回文档数量
        'fetch_k': 20,  # 指定MMR候选集的数量
        'lambda_mult': 0.5,  # 指定MMR系数,兼顾结果的相关性与多样性
    }
).configurable_fields(
    # 指定可配置字段
    search_type=ConfigurableField(id='search_type'),
    search_kwargs=ConfigurableField(id='search_kwargs')
)
result = retriever.invoke('骆驼祥子和刘四爷是什么关系？')
print(result)
